#!/usr/bin/python
import itertools
import os
import re
import string

FLAG_CNVR = 0x01
FLAG_GEXP = 0x02
FLAG_METH = 0x04
FLAG_MIRN = 0x08
FLAG_LEVEL_1 = 0x10
FLAG_LEVEL_2 = 0x20
FLAG_LEVEL_3 = 0x40
FLAG_LEVEL_4 = 0x80

EXCLUDE_MIMETYPES = ["CEL","idat","wig"]

# Todo: This system was never supposed to scale to be inclusive of platform types. Define a better data structure for this.
dataTypeDict = {}
dataTypeDict["Genome_Wide_SNP_6"]   = [ "N", 0x01, 0x100 ]
dataTypeDict["CGH-1x1M_G4447A"]     = [ "N", 0x01, 0x200 ]
dataTypeDict["Human1MDuo"]          = [ "N", 0x01, 0x400 ]
dataTypeDict["HG-CGH-415K_G4124A"]  = [ "N", 0x01, 0x800 ]
dataTypeDict["HG-CGH-244A"]         = [ "N", 0x01, 0x1000 ]
dataTypeDict["HumanHap550"]         = [ "N", 0x01, 0x2000 ]
dataTypeDict["HT_HG-U133A"]         = [ "N", 0x02, 0x4000 ]
dataTypeDict["AgilentG4502A_07_1"]  = [ "N", 0x02, 0x8000 ]
dataTypeDict["AgilentG4502A_07_2"]  = [ "N", 0x02, 0x10000 ]
dataTypeDict["AgilentG4502A_07_3"]  = [ "N", 0x02, 0x20000 ]
dataTypeDict["HuEx-1_0-st-v2"]      = [ "N", 0x02, 0x40000 ]
dataTypeDict["IlluminaGA_RNASeq"]   = [ "N", 0x02, 0x80000 ]
dataTypeDict["HumanMethylation27"]  = [ "N", 0x04, 0x100000 ]
dataTypeDict["H-miRNA_8x15K"]       = [ "N", 0x08, 0x200000 ]
dataTypeDict["IlluminaGA_miRNASeq"] = [ "N", 0x08, 0x400000 ]

# Cancer lookup table, built from the intersection of these two data sets:
#  http://tcga-data.nci.nih.gov/datareports/codeTablesReport.htm?codeTable=Tissue%20Source%20Site
#  http://tcga-data.nci.nih.gov/datareports/codeTablesReport.htm?codeTable=Tissue%20Source%20Site
cancerTypeDict = {}
cancerTypeDict['01'] = "OV"
cancerTypeDict['02'] = "GBM"
cancerTypeDict['03'] = "LUSC"
cancerTypeDict['04'] = "OV"
cancerTypeDict['05'] = "LUAD"
cancerTypeDict['06'] = "GBM"
cancerTypeDict['07'] = "Control"
cancerTypeDict['08'] = "GBM"
cancerTypeDict['09'] = "OV"
cancerTypeDict['10'] = "OV"
cancerTypeDict['11'] = "LUSC"
cancerTypeDict['12'] = "GBM"
cancerTypeDict['13'] = "OV"
cancerTypeDict['14'] = "GBM"
cancerTypeDict['15'] = "GBM"
cancerTypeDict['16'] = "GBM"
cancerTypeDict['17'] = "LUAD"
cancerTypeDict['18'] = "LUSC"
cancerTypeDict['19'] = "GBM"
cancerTypeDict['20'] = "OV"
cancerTypeDict['21'] = "LUSC"
cancerTypeDict['22'] = "LUSC"
cancerTypeDict['23'] = "OV"
cancerTypeDict['24'] = "OV"
cancerTypeDict['25'] = "OV"
cancerTypeDict['26'] = "GBM"
cancerTypeDict['27'] = "GBM"
cancerTypeDict['28'] = "GBM"
cancerTypeDict['29'] = "OV"
cancerTypeDict['30'] = "OV"
cancerTypeDict['31'] = "OV"
cancerTypeDict['32'] = "GBM"
cancerTypeDict['33'] = "LUSC"
cancerTypeDict['34'] = "LUSC"
cancerTypeDict['35'] = "LUAD"
cancerTypeDict['36'] = "OV"
cancerTypeDict['37'] = "LUSC"
cancerTypeDict['38'] = "LUAD"
cancerTypeDict['39'] = "LUSC"
cancerTypeDict['40'] = "LUAD"
cancerTypeDict['41'] = "GBM"
cancerTypeDict['42'] = "OV"
cancerTypeDict['43'] = "LUSC"
cancerTypeDict['44'] = "LUAD"
cancerTypeDict['45'] = "GBM"
cancerTypeDict['46'] = "LUSC"
cancerTypeDict['47'] = "OV"
cancerTypeDict['48'] = "LUAD"
cancerTypeDict['49'] = "LUAD"
cancerTypeDict['50'] = "LUAD"
cancerTypeDict['51'] = "LUSC"
cancerTypeDict['55'] = "LUAD"
cancerTypeDict['56'] = "LUSC"
cancerTypeDict['57'] = "OV"
cancerTypeDict['59'] = "OV"
cancerTypeDict['60'] = "LUSC"
cancerTypeDict['61'] = "OV"
cancerTypeDict['63'] = "LUSC"
cancerTypeDict['64'] = "LUAD"
cancerTypeDict['65'] = "GBM"
cancerTypeDict['66'] = "LUSC"
cancerTypeDict['67'] = "LUAD"
cancerTypeDict['70'] = "LUSC"
cancerTypeDict['71'] = "LUAD"
cancerTypeDict['72'] = "OV"
cancerTypeDict['73'] = "LUAD"
cancerTypeDict['74'] = "GBM"
cancerTypeDict['75'] = "LUAD"
cancerTypeDict['76'] = "GBM"
cancerTypeDict['77'] = "LUSC"
cancerTypeDict['78'] = "LUAD"
cancerTypeDict['79'] = "LUSC"
cancerTypeDict['80'] = "LUAD"
cancerTypeDict['81'] = "GBM"
cancerTypeDict['82'] = "LUSC"
cancerTypeDict['83'] = "LUAD"
cancerTypeDict['84'] = "GBM"
cancerTypeDict['85'] = "LUSC"
cancerTypeDict['86'] = "LUAD"
cancerTypeDict['87'] = "GBM"
cancerTypeDict['88'] = "LUSC"
cancerTypeDict['89'] = "LUAD"
cancerTypeDict['90'] = "LUSC"
cancerTypeDict['91'] = "LUAD"
cancerTypeDict['92'] = "LUSC"
cancerTypeDict['93'] = "LUAD"
cancerTypeDict['94'] = "LUSC"
cancerTypeDict['95'] = "LUAD"
cancerTypeDict['A1'] = "BRCA"
cancerTypeDict['A2'] = "BRCA"
cancerTypeDict['A3'] = "KIRC"
cancerTypeDict['A4'] = "KIRP"
cancerTypeDict['A5'] = "UCEC"
cancerTypeDict['A6'] = "COAD"
cancerTypeDict['A7'] = "BRCA"
cancerTypeDict['A8'] = "BRCA"
cancerTypeDict['AA'] = "COAD"
cancerTypeDict['AB'] = "LAML"
cancerTypeDict['AC'] = "BRCA"
cancerTypeDict['AD'] = "COAD"
cancerTypeDict['AE'] = "COAD"
cancerTypeDict['AF'] = "READ"
cancerTypeDict['AG'] = "READ"
cancerTypeDict['AH'] = "READ"
cancerTypeDict['AI'] = "READ"
cancerTypeDict['AJ'] = "UCEC"
cancerTypeDict['AK'] = "KIRC"
cancerTypeDict['AL'] = "KIRP"
cancerTypeDict['AM'] = "COAD"
cancerTypeDict['AN'] = "BRCA"
cancerTypeDict['AO'] = "BRCA"
cancerTypeDict['AP'] = "UCEC"
cancerTypeDict['AQ'] = "BRCA"
cancerTypeDict['AR'] = "BRCA"
cancerTypeDict['AS'] = "KIRC"
cancerTypeDict['AT'] = "KIRP"
cancerTypeDict['AU'] = "COAD"
cancerTypeDict['AV'] = "Control"
cancerTypeDict['AW'] = "UCEC"
cancerTypeDict['AX'] = "UCEC"
cancerTypeDict['AY'] = "COAD"
cancerTypeDict['AZ'] = "COAD"
cancerTypeDict['B0'] = "KIRC"
cancerTypeDict['B1'] = "KIRP"
cancerTypeDict['B2'] = "KIRC"
cancerTypeDict['B3'] = "KIRP"
cancerTypeDict['B4'] = "KIRC"
cancerTypeDict['B5'] = "UCEC"
cancerTypeDict['B6'] = "BRCA"
cancerTypeDict['B7'] = "STAD"
cancerTypeDict['B8'] = "KIRC"
cancerTypeDict['B9'] = "KIRP"
cancerTypeDict['BA'] = "HNSC"
cancerTypeDict['BB'] = "HNSC"
cancerTypeDict['BC'] = "LIHC"
cancerTypeDict['BD'] = "LIHC"
cancerTypeDict['BE'] = "HNSC"
cancerTypeDict['BF'] = "SKCM"
cancerTypeDict['BG'] = "UCEC"
cancerTypeDict['BH'] = "BRCA"
cancerTypeDict['BI'] = "CESC"
cancerTypeDict['BJ'] = "THCA"
cancerTypeDict['BK'] = "UCEC"
cancerTypeDict['BL'] = "BLCA"
cancerTypeDict['BM'] = "READ"
cancerTypeDict['BP'] = "KIRC"
cancerTypeDict['BQ'] = "KIRP"
cancerTypeDict['BR'] = "STAD"
cancerTypeDict['BS'] = "UCEC"
cancerTypeDict['BT'] = "BLCA"
cancerTypeDict['BV'] = "LGG"
cancerTypeDict['BW'] = "LIHC"
cancerTypeDict['BX'] = "STAD"
cancerTypeDict['BY'] = "HNSC"
cancerTypeDict['BZ'] = "THCA"
cancerTypeDict['C2'] = "DLBC"
cancerTypeDict['C3'] = "PRAD"
cancerTypeDict['C4'] = "BLCA"
cancerTypeDict['C5'] = "CESC"
cancerTypeDict['C6'] = "BRCA"
cancerTypeDict['C7'] = "BRCA"
cancerTypeDict['C8'] = "BRCA"
cancerTypeDict['CA'] = "COAD"
cancerTypeDict['CB'] = "KIRC"
cancerTypeDict['CC'] = "LIHC"
cancerTypeDict['CD'] = "STAD"
cancerTypeDict['CE'] = "THCA"
cancerTypeDict['CF'] = "BLCA"
cancerTypeDict['CG'] = "STAD"
cancerTypeDict['CH'] = "PRAD"
cancerTypeDict['CI'] = "READ"
cancerTypeDict['CJ'] = "KIRC"
cancerTypeDict['CK'] = "COAD"
cancerTypeDict['CL'] = "READ"
cancerTypeDict['CM'] = "COAD"
cancerTypeDict['CN'] = "HNSC"
cancerTypeDict['CP'] = "LGG"
cancerTypeDict['CQ'] = "HNSC"
cancerTypeDict['CR'] = "HNSC"
cancerTypeDict['CS'] = "LGG"
cancerTypeDict['CU'] = "BLCA"
cancerTypeDict['CV'] = "HNSC"
cancerTypeDict['CW'] = "KIRC"
cancerTypeDict['CX'] = "HNSC"
cancerTypeDict['CY'] = "COAD"
cancerTypeDict['CZ'] = "KIRC"
cancerTypeDict['D1'] = "UCEC"
cancerTypeDict['D2'] = "UCEC"
cancerTypeDict['D3'] = "SKCM"
cancerTypeDict['D4'] = "LAML"
cancerTypeDict['D5'] = "COAD"
cancerTypeDict['D6'] = "HNSC"
cancerTypeDict['D7'] = "STAD"
cancerTypeDict['D8'] = "BRCA"
cancerTypeDict['D9'] = "SKCM"
cancerTypeDict['DA'] = "SKCM"
cancerTypeDict['DB'] = "LGG"
cancerTypeDict['DC'] = "READ"
cancerTypeDict['DD'] = "LIHC"
cancerTypeDict['DE'] = "THCA"
cancerTypeDict['DF'] = "UCEC"
cancerTypeDict['DG'] = "CESC"
cancerTypeDict['DH'] = "LGG"
cancerTypeDict['DI'] = "UCEC"
cancerTypeDict['DJ'] = "THCA"
cancerTypeDict['DK'] = "BLCA"
cancerTypeDict['DM'] = "COAD"
cancerTypeDict['DN'] = "LAML"
cancerTypeDict['DO'] = "THCA"
cancerTypeDict['DQ'] = "HNSC"
cancerTypeDict['DR'] = "CESC"
cancerTypeDict['DS'] = "CESC"
cancerTypeDict['DT'] = "READ"
cancerTypeDict['DU'] = "LGG"
cancerTypeDict['DV'] = "KIRC"
cancerTypeDict['DW'] = "KIRP"
cancerTypeDict['DX'] = "SALD"
cancerTypeDict['DY'] = "READ"
cancerTypeDict['DZ'] = "KIRP"
cancerTypeDict['E1'] = "LGG"
cancerTypeDict['E2'] = "BRCA"
cancerTypeDict['E3'] = "THCA"
cancerTypeDict['E5'] = "BLCA"
cancerTypeDict['E6'] = "UCEC"
cancerTypeDict['E7'] = "BLCA"
cancerTypeDict['E8'] = "THCA"
cancerTypeDict['E9'] = "BRCA"
cancerTypeDict['EA'] = "CESC"
cancerTypeDict['EB'] = "SKCM"
cancerTypeDict['EC'] = "UCEC"
cancerTypeDict['ED'] = "LIHC"
cancerTypeDict['EE'] = "SKCM"
cancerTypeDict['EF'] = "READ"
cancerTypeDict['EG'] = "READ"
cancerTypeDict['EH'] = "READ"
cancerTypeDict['EI'] = "READ"
cancerTypeDict['EJ'] = "PRAD"
cancerTypeDict['EK'] = "CESC"
cancerTypeDict['EL'] = "THCA"
cancerTypeDict['EM'] = "THCA"
cancerTypeDict['EN'] = "CESC"
cancerTypeDict['EO'] = "UCEC"
cancerTypeDict['EP'] = "LIHC"
cancerTypeDict['EQ'] = "STAD"
cancerTypeDict['ER'] = "SKCM"
cancerTypeDict['ES'] = "LIHC"
cancerTypeDict['ET'] = "THCA"
cancerTypeDict['EU'] = "KIRC"
cancerTypeDict['EV'] = "KIRP"
cancerTypeDict['EW'] = "BRCA"
cancerTypeDict['EX'] = "CESC"
cancerTypeDict['EY'] = "UCEC"
cancerTypeDict['EZ'] = "LGG"
cancerTypeDict['F1'] = "STAD"
cancerTypeDict['F2'] = "PAAD"
cancerTypeDict['F4'] = "COAD"
cancerTypeDict['F5'] = "READ"
cancerTypeDict['F6'] = "LGG"
cancerTypeDict['F7'] = "HNSC"
cancerTypeDict['F8'] = "KIRC"
cancerTypeDict['F9'] = "KIRP"
cancerTypeDict['FA'] = "DLBC"
cancerTypeDict['FB'] = "PAAD"
cancerTypeDict['FC'] = "PRAD"
cancerTypeDict['FD'] = "BLCA"
cancerTypeDict['FE'] = "THCA"
cancerTypeDict['FF'] = "DLBC"
cancerTypeDict['FG'] = "LGG"
cancerTypeDict['FH'] = "THCA"
cancerTypeDict['FI'] = "UCEC"
cancerTypeDict['FJ'] = "BLCA"
cancerTypeDict['FK'] = "THCA"
cancerTypeDict['FL'] = "UCEC"
cancerTypeDict['FM'] = "DLBC"
cancerTypeDict['FN'] = "LGG"
cancerTypeDict['FP'] = "STAD"
cancerTypeDict['FQ'] = "PAAD"
cancerTypeDict['FS'] = "SKCM"
cancerTypeDict['FT'] = "BLCA"
cancerTypeDict['FU'] = "CESC"
cancerTypeDict['FV'] = "LIHC"
cancerTypeDict['FW'] = "SKCM"
cancerTypeDict['FX'] = "SALD"
cancerTypeDict['FY'] = "THCA"
cancerTypeDict['FZ'] = "PAAD"
cancerTypeDict['G1'] = "LAML"
cancerTypeDict['G2'] = "BLCA"
cancerTypeDict['G3'] = "LIHC"
cancerTypeDict['G4'] = "COAD"
cancerTypeDict['G5'] = "READ"
cancerTypeDict['G6'] = "KIRC"
cancerTypeDict['G7'] = "KIRP"
cancerTypeDict['G8'] = "DLBC"
cancerTypeDict['G9'] = "PRAD"
cancerTypeDict['GA'] = "LNNH"
cancerTypeDict['GB'] = "LNNH"
cancerTypeDict['GD'] = "BLCA"
cancerTypeDict['GE'] = "THCA"
cancerTypeDict['GF'] = "SKCM"
cancerTypeDict['GG'] = "UCEC"
cancerTypeDict['GH'] = "CESC"
cancerTypeDict['GI'] = "BRCA"
cancerTypeDict['GJ'] = "LIHC"
cancerTypeDict['GK'] = "KIRC"
cancerTypeDict['GL'] = "KIRP"
cancerTypeDict['GM'] = "BRCA"
cancerTypeDict['GN'] = "SKCM"
cancerTypeDict['GP'] = "LAML"
cancerTypeDict['GQ'] = "LNNH"
cancerTypeDict['GR'] = "DLBC"
cancerTypeDict['GS'] = "DLBC"
cancerTypeDict['GT'] = "COAD"
cancerTypeDict['GU'] = "BLCA"
cancerTypeDict['GV'] = "BLCA"
cancerTypeDict['GW'] = "STAD"
cancerTypeDict['GX'] = "PRAD"
cancerTypeDict['GY'] = "KIRP"
cancerTypeDict['GZ'] = "DLBC"
cancerTypeDict['H1'] = "STAD"
cancerTypeDict['H2'] = "THCA"
cancerTypeDict['H3'] = "DLBC"
cancerTypeDict['H4'] = "BLCA"
cancerTypeDict['H5'] = "UCEC"
cancerTypeDict['H6'] = "PAAD"
cancerTypeDict['H7'] = "HNSC"
cancerTypeDict['H8'] = "PAAD"
cancerTypeDict['H9'] = "PRAD"


barcodes = {}

classified = 0
unclassified = 0

def doParse(root, dir, path, level_def):

	global classified,unclassified
	# Establish the level of this data

	# Path into each of the contained files, reading their 
	#  first line. We use this to determine the type of file.
	found = False
	for i in dataTypeDict.keys():
		if i in root:
			for file in path:
				if '.' in file and file.rsplit('.',1)[1] in EXCLUDE_MIMETYPES:
					continue
				
				file = os.path.join(root,file)
				print "Reading: " + file
				samples = []
				if os.path.isfile(file):
					#contents = open(file,'r').read()
					reader = open(file,'r')
					line = '\n'
					count = 0
					top = 3
					if 'mage-tab' in root:
						top = 1 << 32;
					while count < top and line != '':
						count += 1
						line = reader.readline()
					
						m = re.search('(TCGA-[0-9A-Za-z]{2}-[0-9A-Za-z]{4}-[0-9A-Za-z]{3}[0-9A-Za-z-]+)[^0-9A-Za-z-]',line)
						if m != None: 
							for sample in m.groups():
								samples.append(sample)
				
				m = re.search('(TCGA-[0-9A-Za-z-]+)[^0-9A-Za-z-]',file)
				if m != None: 
					for sample in m.groups():
						samples.append(sample)

				for sample in samples:
					current = 0
					sample = sample[:16]
					print sample
					if sample in barcodes:
						current = barcodes[sample]
					barcodes[sample] = current | dataTypeDict[i][1] | level_def | dataTypeDict[i][2]

			found = True
			classified += 1
			break
	if not found:
		if 'Level_1' not in root and 'mage-tab' not in root and 'gdacbroad' not in root:
			raise NameError('Could not find definition for ' + root)
			unclassified += 1
	#for file in path:
	#	print open(os.path.join(root,file),'r').readline()
	#	print open(os.path.join(root,file),'r').readline()
	#	print open(os.path.join(root,file),'r').readline()


# Support Firehose's reparsing of platform names by adding compatibility definitions
for i in dataTypeDict.keys():
	dataTypeDict[i.lower().replace('-','_')] = dataTypeDict[i]

print dataTypeDict.keys()

for root, dir, path in os.walk('/titan/cancerregulome7/TCGA/repositories/dcc-mirror'):
	#print root,dir,path
	root_lower = root.lower()
	if 'level_' in root_lower or 'mage-tab' in root_lower:
		level_def = -1
		if 'level_1' in root_lower:
			level_def = 0x10
		if 'level_2' in root_lower:
			level_def = 0x20
		if 'level_3' in root_lower:
			level_def = 0x40
		if level_def == -1:
			level_def = 0x00
		doParse(root,dir,path,level_def)

for root, dir, path in os.walk('/titan/cancerregulome7/TCGA/gdacbroad'):
	#print root,dir,path
	root_lower = root.lower()
	if 'level_' in root_lower:
		level_def = 0x80	# All data at this location should be Level 4, regardless of what it claims to be.
		doParse(root,dir,path,level_def)
					
if unclassified > 0:
	print 'WARNING: Some groups were not classified. Parsing may not \
		be robust!'

files = {}
combinations = {}
cancerTypeDict[''] = 'ALL' # Catch-all entry
for cancerId in cancerTypeDict.values():
        if cancerId not in files:
		combinations[cancerId] = {}
                files[cancerId] = open(cancerId + '.json','w')
                files[cancerId].write('{\n')

writer = open('out.tsv','w')
for i in barcodes.keys():
	# In place update to flatten level and platform data	
	barcodes[i] &= 0xf

	writer.write(i + '\t' + str(barcodes[i]) + "\n")
	sampleSite = i[5:7]
	if barcodes[i] not in combinations[cancerTypeDict[sampleSite]]:
		combinations[cancerTypeDict[sampleSite]][barcodes[i]] = 0
	combinations[cancerTypeDict[sampleSite]][barcodes[i]] += 1

	if barcodes[i] not in combinations['ALL']:
		combinations['ALL'][barcodes[i]] = 0
	combinations['ALL'][barcodes[i]] += 1
writer.close()

for i in combinations.keys():
	for j in combinations[i].keys():
		props = []
		if j & FLAG_CNVR:
			props.append("Copy Number Variation")
		if j & FLAG_GEXP:
			props.append("Gene Expression")
		if j & FLAG_METH:
			props.append("Methylation")
		if j & FLAG_MIRN:
			props.append("MiRNA")
		'''
		if j & FLAG_LEVEL_1:
			props.append("Level 1")
		if j & FLAG_LEVEL_2:
			props.append("Level 2")
		if j & FLAG_LEVEL_3:
			props.append("Level 3")
		if j & FLAG_LEVEL_4:
			props.append("Level 4")
		if j & 0x100:
			props.append("SNP6") #Genome_Wide_SNP_6")
		if j & 0x200:
			props.append("1x1M") #CGH-1x1M_G4447A")
		if j & 0x400:
			props.append("H1MD") #Human1MDuo")
		if j & 0x800:
			props.append("415K") #HG-CGH-415K_G4124A")
		if j & 0x1000:
			props.append("244A") #HG-CGH-244A")
		if j & 0x2000:
			props.append("H550") #HumanHap550")
		if j & 0x4000:
			props.append("U133") #HT_HG-U133A")
		if j & 0x8000:
			props.append("AG71") #AgilentG4502A_07_1")
		if j & 0x10000:
			props.append("AG72") #AgilentG4502A_07_2")
		if j & 0x20000:
			props.append("AG73") #AgilentG4502A_07_3")
		if j & 0x40000:
			props.append("HEv2") #HuEx-1_0-st-v2")
		if j & 0x80000:
			props.append("IRSq") #IlluminaGA_RNASeq")
		if j & 0x100000:
			props.append("HM27") #HumanMethylation27")
		if j & 0x200000:
			props.append("8x15") #H-miRNA_8x15K")
		if j & 0x400000:
			props.append("ImRS") #IlluminaGA_miRNASeq")
		'''

                files[i].write('"' + string.join(props,', ') + '": ' + str(combinations[i][j]) + ',\n')

for i in files.keys():
	if combinations[i] == {}:
		files[i].write('"No data": ' + str(1 << 24) + '\n');
	files[i].write('}')
	files[i].close()
